import torch
import torchvision
from torch import nn
from d2l import torch as d2l
import os

normalize = torchvision.transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
train_augs = torchvision.transforms.Compose([
    torchvision.transforms.RandomResizedCrop(224),
    torchvision.transforms.RandomHorizontalFlip(),
    torchvision.transforms.ToTensor(),
    normalize
])
test_augs = torchvision.transforms.Compose([
    torchvision.transforms.Resize(256),
    torchvision.transforms.CenterCrop(224),
    torchvision.transforms.ToTensor(),
    normalize
])

d2l.DATA_HUB['hotdog'] = (d2l.DATA_URL+'hotdog.zip',
                          'fba480ffa8aa7e0febbb511d181409f899b9baa5')
data_dir = d2l.download_extract('hotdog')

def data_test_show():
    train_imgs = torchvision.datasets.ImageFolder(os.path.join(data_dir,'train'))
    test_imgs = torchvision.datasets.ImageFolder(os.path.join(data_dir,'test'))

    hotdogs = [train_imgs[i][0] for i in range(8)]
    not_hotdogs = [train_imgs[-i-1][0] for i in range(8)]
    d2l.show_images(hotdogs+not_hotdogs,2,8,scale=1.4)
    d2l.plt.show()


pretrain_net = torchvision.models.resnet18(pretrained=True)
pretrain_net.fc = nn.Linear(pretrain_net.fc.in_features,2)
nn.init.xavier_uniform_(pretrain_net.fc.weight)

def train_fine_tuning(net,lr,batch_size=128,num_epochs=10,param_group=True):
    train_imgs = torchvision.datasets.ImageFolder(os.path.join(data_dir,'train'),transform=train_augs)
    test_imgs = torchvision.datasets.ImageFolder(os.path.join(data_dir,'test'),transform=test_augs)
    train_iter = torch.utils.data.DataLoader(train_imgs,batch_size=batch_size,shuffle=True)
    test_iter = torch.utils.data.DataLoader(test_imgs,batch_size=batch_size)
    devices = d2l.try_all_gpus()
    loss = nn.CrossEntropyLoss(reduction='none')
    if param_group:
        params_1x = [
            param for name,param in net.named_parameters()
                if name not in ["fc.weight","fc.bias"]]
        trainer = torch.optim.SGD([{'params':params_1x},
                                   {'params':net.fc.parameters(),'lr':lr*10}],
                                   lr=lr,weight_decay=0.001)
    else:
        trainer = torch.optim.SGD(net.parameters(),lr=lr,weight_decay=0.001)
    d2l.train_ch13(net,train_iter,test_iter,loss,trainer,num_epochs,devices)

train_fine_tuning(pretrain_net,5e-5)

        


